from tensorflow.keras.preprocessing.image import ImageDataGenerator

"""
image generator: tf.keras.preprocessing.image.ImageDataGenerator
parameters:
- rescale: rescale pixels from range 0 ~255 ro range 0~1 by set 1/255

import pics from dir: tf.keras.prepocessing.image.ImageDataGenerator.flow_from_directory
parameters:
- directory: pics path
- target_size: transform heights*weights to standard scale, default = (256, 256)
- batch_size: default = 32
- class_mode: default = 'categorical'
              example: 'sparse': ['paper', 'rock', 'scissors'] --> [0, 1, 2]
                       'categorical': ['paper', 'rock', 'scissors'] --> [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
                       'input': ['paper', 'rock', 'scissors'] --> ['paper', 'rock', 'scissors']
                       'None': no label return  
"""


def train_val_generator(data_dir, target_size, batch_size, class_mode=None, subset='training'):
    train_val_datagen = ImageDataGenerator(rescale=1./255., validation_split=0.2)
    return train_val_datagen.flow_from_directory(
        directory=data_dir,
        target_size=target_size,
        batch_size=batch_size,
        class_mode=class_mode,
        subset=subset
    )


def test_generator(data_dir, target_size, batch_size, class_mode=None):
    test_datagen = ImageDataGenerator(rescale=1./255.)
    return test_datagen.flow_from_directory(
        directory=data_dir,
        target_size=target_size,
        batch_size=batch_size,
        class_mode=class_mode
    )


def pred_generator(data_dir, target_size, batch_size, class_mode=None):
    pred_datagen = ImageDataGenerator(rescale=1./255.)
    return pred_datagen.flow_from_directory(
        directory=data_dir,
        target_size=target_size,
        batch_size=batch_size,
        class_mode=class_mode
    )
